Pytorch Selfattention Tutorial